# -*-coding:utf-8-*-

from typing import Union, List
import  numpy as np

def AAINDEX(sequence: str) -> List:
    '''
    AAindex is a database of numerical indices representing various physicochemical and biochemical properties of amino acids and pairs of amino acids. AAindex consists of three sections now: AAindex1 for the amino acid index of 20 numerical values, AAindex2 for the amino acid mutation matrix and AAindex3 for the statistical protein contact potentials. All data are derived from published literature.
    这里只用了index1.txt
    :param sequence:string 序列
    :return: List
    '''
    order = "ARNDCQEGHILKMFPSTWYV-"
    with open('data/AAindex.txt') as f:
        records = f.readlines()[1:]

    AAindex = []
    AAindexName = []
    for i in records:
        AAindex.append(i.rstrip().split()[1:] if i.rstrip() != '' else None)
        AAindexName.append(i.rstrip().split()[0] if i.rstrip() != '' else None)
    index = {}
    for i in range(len(order)):
        index[order[i]] = i

    encodings = []
    for aa in sequence:
        if aa == '-':
            for j in AAindex:
                encodings.append(0)
            continue
        if aa not in index:
            for j in AAindex:
                encodings.append(0)
            continue
        # print(aa)
        for j in AAindex:
            encodings.append(j[index[aa]])
    return encodings

def one_AAINDEX(data, windows=15):
    data = data
    length = len(data)
    data_X = np.zeros((length, (2*windows+1), 531))
    data_Y = []
    for i in range(length):
        x = data[i].split(',')
        data_Y.append(int(x[4]))
        aaindex = np.array(AAINDEX(x[5]))
        data_X[i] = list(aaindex.reshape((2*windows+1),531))
    data_Y = np.array(data_Y)
    data_X = np.array(data_X)
    return data_X, data_Y

if __name__ == '__main__':
    f_r_train = open('hit-40/totalsum.csv')
    train_data = f_r_train.readlines()
    result = one_AAINDEX(train_data)
    print(result.shape)
